from khronos import *

class ProductType(object):
    def __init__(self, name, 
                 ptime_mean, 
                 ptime_stddev):
        self.name = name
        self.mu = ptime_mean
        self.sigma = ptime_stddev
        
    def __repr__(self):
        return self.name
        
class Order(object):
    def __init__(self, type, qty, due):
        self.type = type
        self.qty = qty
        self.due = due
        
    def __repr__(self):
        return "Order(%d x %s, due %s)" % (self.qty, self.type, self.due)
        
    def ptime(self, rng):
        return self.qty * rng.gauss(self.type.mu, self.type.sigma)
        
quantities = 10000, 20000, 30000
delays = 1*WEEK, 2*WEEK, 3*WEEK
types = (ProductType("red", 5, 0.5), 
         ProductType("green", 6, 0.25), 
         ProductType("blue", 9,  1.1))

class Source(Atomic):
    @chain
    def initializer(self):
        global quantities, delays, types
        
        rng = self.get_rng()
        queue = self.find(">chassis.queue")
        rrate = lambda: rng.expovariate(10) * WEEK
        
        yield rng.random() * rrate()
        while True:
            type = rng.choice(types)
            qty = rng.choice(quantities)
            now = self.get_timedelta()
            delay = rng.choice(delays)
            duedate = now + delay
            order = Order(type, qty, duedate)
            self.send(order, queue)
            yield rrate() # time until next order
            
class Sink(Atomic):
    def initializer(self):
        self.lateness = Tally()
        self.avg_lateness = TSeries(timescale=DAY, storing=True)
        
    def receive(self, order, _):
        now = self.get_timedelta()
        lateness = max(0, (now-order.due) / DAY)
        self.lateness.register(lateness)
        self.avg_lateness.register(self.lateness.mean(), now)
        
    def finalizer(self):
        self.avg_lateness.run_chart(show_mean=False)
        
def FIFO(orders, machine):
    return [-x for x in xrange(len(orders))]
    
def manager(orders, machine):
    return [-order.due for order in orders]
    
def foreman(orders, machine):
    return [1 if order.type is machine.setup else 0 for order in orders]
    
def operator(orders, machine):
    return [-(order.qty * order.type.mu) for order in orders]

class Queue(Atomic):
    default_priority = 1.0
    
    def __init__(self, policy=FIFO, *args, **kwargs):
        Atomic.__init__(self, *args, **kwargs)
        self.policy = policy
        self.content = []
        self.reqs = []
        
    def status(self):
        return "%s (%d requests)" % (self.content, len(self.reqs))
        
    def initializer(self):
        self.content = []
        self.reqs = []
        
    def receive(self, order, _):
        if self.reqs:
            self.reqs[0].satisfy(order)
        else:
            self.content.append(order)
            
    def get(self, machine):
        req = Request(deploy=self.reqs.append, 
                      retract=self.reqs.remove)
        if self.content:
            orders = self.content
            scores = self.policy(orders, machine)
            _, order = max(zip(scores, orders))
            orders.remove(order)
            req.satisfy(order)
        return req
        
    def finalizer(self):
        stat = FTable()
        for m in self.get_parent():
            if m is not self:
                stat.merge(Machine.state.stat(m))
        axes = stat.pareto_chart(show=False)
        axes.figure.plotter.set_title("Toy example")
        
class Machine(Atomic):
    state = FTableAttr()
    
    def __init__(self, queue=None, next=None, *args, **kwargs):
        Atomic.__init__(self, *args, **kwargs)
        Machine.state.initialize(self)
        self.queue = queue
        self.next = next
        self.order = None
        
    def status(self):
        return "[%s] %s" % (self.state, self.order)
        
    @chain
    def initializer(self):
        Machine.state.reset(self)
        rng = self.get_rng()
        self.setup = None
        self.order = None
        
        while True:
            # request an order from the queue
            self.state = "SB"
            order = (yield self.queue.get(self)).result
            self.order = order
            
            # perform a conversion if necessary
            if order.type is not self.setup:
                self.state = "CO"
                yield HOUR
                self.setup = order.type
                
            # process
            self.state = "PR"
            yield order.ptime(rng)
            
            # forward order
            self.send(order, self.next)
            
def build_model():
    m = Composite("jobshop")
    Source("entry", parent=m)
    next = Sink("exit", parent=m)
    for op in "trim", "body", "chassis":
        op = Composite(op, parent=m)
        queue = Queue(name="queue", parent=op)
        for x in xrange(2):
            Machine(queue, next, name="m%d" % (x,), parent=op)
        next = queue
    return m
    
def build():
    k = Khronos(build_model())
    return k, k.model
    
def main():
    k, m = build()
    k.stack.trace.disable()
    k.single_run(10 * WEEK)
    
k, m = build()
